#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   train.py
@Author :   Song
@Time   :   2022/2/28 21:48
@Contact:   songjian@westlake.edu.cn
@intro  : 
'''
import numpy as np
import einops
import matplotlib.pyplot as plt
import torch.optim
import torch.utils.data
import torch.nn.functional as F
import torch.nn.utils.rnn as rnn_utils
from pathlib import Path
import predifine
import AlphaNovo_models
from AlphaNovo_dataloader import TrainEval_Dataset

try:
    profile
except:
    profile = lambda x: x

def eval_one_epoch(dataloader, model, epoch):
    model.eval()
    batch_match_num_v, batch_match_num_len_bias_v = [], []
    for batch_idx, (batch_spectra, batch_pr_charge, batch_num) in enumerate(dataloader):

        batch_spectra = batch_spectra.float().to(device)
        batch_pr_charge = batch_pr_charge.long().to(device)
        batch_num = batch_num.long().to(device)

        # forward
        with torch.no_grad():
            # forward
            pred_num = model(batch_spectra, batch_pr_charge)

            # for pred len
            pred_num = torch.softmax(pred_num, dim=1)
            pred_num = torch.argmax(pred_num, dim=1)
            batch_match_num = ((pred_num <= batch_num + 1) & (pred_num >= batch_num - 1)).sum()
            batch_acc = batch_match_num / len(batch_num)
            batch_match_num_v.append(batch_match_num.item())

            if batch_idx % 100 == 0:
                print('Eval epoch: [{}], Batch: [{}/{}], '
                      'batch acc: {:.3f}'.format(
                    epoch, batch_idx, len(dataloader),
                    batch_acc.item()))

    epoch_acc = sum(batch_match_num_v) / len(dataloader.dataset)
    print('Eval epoch {}, epoch_acc: {:.3f}, '.format(
        epoch, epoch_acc))

    return epoch_acc


def my_collate(items):
    batch_spectra, batch_pr_charge, batch_num = zip(*items)

    batch_spectra = rnn_utils.pad_sequence(batch_spectra, batch_first=True)
    batch_pr_charge = torch.tensor(batch_pr_charge)
    batch_num = torch.tensor(batch_num)

    return batch_spectra, batch_pr_charge, batch_num


if __name__ == '__main__':
    # train and valid
    device = predifine.device
    eval_dataset = TrainEval_Dataset('eval', path=predifine.eval_npz)

    eval_loader = torch.utils.data.DataLoader(eval_dataset,
                                              batch_size=predifine.batch_size,
                                              num_workers=predifine.num_workers,
                                              shuffle=predifine.shuffle,
                                              pin_memory=True,
                                              collate_fn=my_collate)

    # %% model
    model = AlphaNovo_models.AlphaNovo_Model_form(
        dim_model=predifine.dim_model,
        n_head=predifine.n_head,
        dim_feedforward=predifine.dim_feedforward,
        n_layers=predifine.n_layers,
        dropout=predifine.dropout,
        dim_intensity=predifine.dim_intensity,
        max_length=predifine.max_length,
        max_charge=predifine.max_charge
    )
    model.load_state_dict(torch.load('/home/songjian/casanovo/train_output_pred_form/alphanovo_form_epoch_9.pt'))
    model = model.to(device)

    eval_acc = eval_one_epoch(eval_loader, model, 0)
    print(f'eval_acc: {eval_acc:.3f}')